package com.jd.genie.service.impl;


import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.agent.ExecutorAgent;
import com.jd.genie.agent.agent.PlanningAgent;
import com.jd.genie.agent.agent.SummaryAgent;
import com.jd.genie.agent.dto.File;
import com.jd.genie.agent.dto.TaskSummaryResult;
import com.jd.genie.agent.enums.AgentState;
import com.jd.genie.agent.enums.AgentType;
import com.jd.genie.agent.util.ThreadUtil;
import com.jd.genie.config.GenieConfig;
import com.jd.genie.model.req.AgentRequest;
import com.jd.genie.service.AgentHandlerService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

/**
 * 计划-解决处理器实现类
 * 
 * 该类实现了基于计划-解决模式的智能代理处理逻辑：
 * 1. 使用规划代理(PlanningAgent)制定任务计划
 * 2. 使用执行代理(ExecutorAgent)执行具体任务
 * 3. 使用总结代理(SummaryAgent)对任务结果进行总结
 * 4. 支持多任务并发执行和迭代优化
 * 
 * @author Genie Team
 */
@Slf4j
@Component
public class PlanSolveHandlerImpl implements AgentHandlerService {

    /**
     * Genie配置信息，包含最大步数等参数
     */
    @Autowired
    private GenieConfig genieConfig;


    /**
     * 处理智能代理请求的核心方法
     * 
     * 执行流程：
     * 1. 初始化三种代理：规划、执行、总结
     * 2. 进行任务规划
     * 3. 循环执行任务直到完成或达到最大步数
     * 4. 对任务结果进行总结并返回
     * 
     * @param agentContext 代理上下文，包含执行环境和状态信息
     * @param request 代理请求，包含用户查询和配置参数
     * @return 处理结果字符串
     */
    @Override
    public String handle(AgentContext agentContext, AgentRequest request) {

        // 初始化三个核心代理
        PlanningAgent planning = new PlanningAgent(agentContext);  // 规划代理：负责制定任务计划
        ExecutorAgent executor = new ExecutorAgent(agentContext);   // 执行代理：负责执行具体任务
        SummaryAgent summary = new SummaryAgent(agentContext);     // 总结代理：负责总结任务结果
        
        // 为总结代理设置系统提示词，将用户查询嵌入到提示词中
        summary.setSystemPrompt(summary.getSystemPrompt().replace("{{query}}", request.getQuery()));

        // 第一步：进行任务规划，获取初始规划结果
        String planningResult = planning.run(agentContext.getQuery());
        log.info("===== {} ==== 获取到 planningResult: \n {} \n ======= \n ", agentContext.getRequestId(), planningResult);
        // 初始化步骤计数器和最大步数限制
        int stepIdx = 0;
        int maxStepNum = genieConfig.getPlannerMaxSteps();
        
        // 主循环：迭代执行任务直到完成或达到最大步数
        while (stepIdx <= maxStepNum) {
            // 解析规划结果，将任务按<sep>分隔符拆分成多个子任务
            List<String> planningResults = Arrays.stream(planningResult.split("<sep>"))
                    .map(task -> "你的任务是：" + task)  // 为每个任务添加前缀说明
                    .toList();
            
            String executorResult;
            // 清空当前任务产生的文件列表，准备执行新任务
            agentContext.getTaskProductFiles().clear();
            
            if (planningResults.size() == 1) {
                // 单任务执行：直接使用执行代理处理
                executorResult = executor.run(planningResults.get(0));
            } else {
                // 多任务并发执行：使用线程池并行处理多个子任务
                Map<String, String> tmpTaskResult = new ConcurrentHashMap<>();  // 存储各任务执行结果
                CountDownLatch taskCount = ThreadUtil.getCountDownLatch(planningResults.size());  // 同步计数器
                int memoryIndex = executor.getMemory().size();  // 记录当前内存索引位置
                List<ExecutorAgent> slaveExecutors = new ArrayList<>();  // 从执行代理列表
                
                // 为每个子任务创建独立的执行代理并启动线程执行
                for (String task : planningResults) {
                    ExecutorAgent slaveExecutor = new ExecutorAgent(agentContext);
                    // 复制主执行代理的状态和内存到从代理
                    slaveExecutor.setState(executor.getState());
                    slaveExecutor.getMemory().addMessages(executor.getMemory().getMessages());
                    slaveExecutors.add(slaveExecutor);
                    
                    // 异步执行子任务
                    ThreadUtil.execute(() -> {
                        String taskResult = slaveExecutor.run(task);
                        tmpTaskResult.put(task, taskResult);  // 保存任务执行结果
                        taskCount.countDown();  // 完成计数
                    });
                }
                
                // 等待所有子任务完成
                ThreadUtil.await(taskCount);
                
                // 合并所有从执行代理的内存和状态到主执行代理
                for (ExecutorAgent slaveExecutor : slaveExecutors) {
                    // 将从代理新增的内存消息合并到主代理
                    for (int i = memoryIndex; i < slaveExecutor.getMemory().size(); i++) {
                        executor.getMemory().addMessage(slaveExecutor.getMemory().get(i));
                    }
                    // 清理从代理内存，避免内存泄漏
                    slaveExecutor.getMemory().clear();
                    // 更新主代理状态
                    executor.setState(slaveExecutor.getState());
                }
                
                // 合并所有任务执行结果
                executorResult = String.join("\n", tmpTaskResult.values());
            }
            
            // 基于执行结果进行下一轮规划
            planningResult = planning.run(executorResult);
            
            if ("finish".equals(planningResult)) {
                // 任务完成：进行结果总结和文件整理
                TaskSummaryResult result = summary.summaryTaskResult(executor.getMemory().getMessages(), request.getQuery());

                // 构建最终任务结果
                Map<String, Object> taskResult = new HashMap<>();
                taskResult.put("taskSummary", result.getTaskSummary());  // 添加任务总结

                // 处理文件列表：优先使用总结代理提取的文件，否则使用上下文中的产品文件
                if (CollectionUtils.isEmpty(result.getFiles())) {
                    if (!CollectionUtils.isEmpty(agentContext.getProductFiles())) {
                        List<File> fileResponses = agentContext.getProductFiles();
                        // 过滤掉内部搜索结果文件，只保留用户需要的文件
                        fileResponses.removeIf(file -> Objects.nonNull(file) && file.getIsInternalFile());
                        // 反转文件列表，让最新的文件排在前面
                        Collections.reverse(fileResponses);
                        taskResult.put("fileList", fileResponses);
                    }
                } else {
                    taskResult.put("fileList", result.getFiles());
                }

                // 发送最终结果给客户端
                agentContext.getPrinter().send("result", taskResult);

                // 任务成功完成，跳出循环
                break;
            }
            
            // 检查代理状态：如果规划或执行代理进入空闲状态，说明达到最大迭代次数
            if (planning.getState() == AgentState.IDLE || executor.getState() == AgentState.IDLE) {
                agentContext.getPrinter().send("result", "达到最大迭代次数，任务终止。");
                break;
            }
            
            // 检查代理状态：如果规划或执行代理出现错误，终止任务执行
            if (planning.getState() == AgentState.ERROR || executor.getState() == AgentState.ERROR) {
                agentContext.getPrinter().send("result", "任务执行异常，请联系管理员，任务终止。");
                break;
            }
            
            // 增加步骤计数器，继续下一轮迭代
            stepIdx++;
        }

        return "";
    }

    /**
     * 判断当前处理器是否支持处理指定的代理请求
     * 
     * @param agentContext 代理上下文
     * @param request 代理请求
     * @return 如果请求类型为PLAN_SOLVE则返回true，否则返回false
     */
    @Override
    public Boolean support(AgentContext agentContext, AgentRequest request) {
        return AgentType.PLAN_SOLVE.getValue().equals(request.getAgentType());
    }
}

